import numpy as np
import paddle.fluid as fluid

# 创建执行器
import sys

place = fluid.CPUPlace()
exe = fluid.Executor(place)
exe.run(fluid.default_startup_program())

# 保存预测模型路径
save_path = '/home/more/src/paddle/infer_model/'
# 从模型中获取预测程序、输入数据名称列表、分类器
[infer_program, feeded_var_names, target_var] = fluid.io.load_inference_model(dirname=save_path, executor=exe)


# 获取数据
def get_data(sentence):
    # 读取数据字典
    with open('/home/more/src/paddle/datasets/dict_txt.txt', 'r', encoding='utf-8') as f_data:
        dict_txt = eval(f_data.readlines()[0])
    dict_txt = dict(dict_txt)
    # 把字符串数据转换成列表数据
    keys = dict_txt.keys()
    data = []
    for s in sentence:
        # 判断是否存在未知字符
        if not s in keys:
            s = '<unk>'
        data.append(np.int64(dict_txt[s]))
    return data


data = []
# 获取图片数据
data1 = get_data(sys.argv[1])
# data1 = get_data('借助机器人视觉技术的灵活自动化解决方案和确保安全、能够自由接触的机器人应用——这就是已经到来的未来。'
#                  'SICK的传感器解决方案让这幅未来图景成为可能。这是一幅人与机器携手工作的图景——正如SICK与其客户在平等合作中探寻解决方案。'
#                  '机器人视觉——这种基于图像的光学系统让机器人成为有“眼睛”的伙伴，使其能够识别出在哪些位置存在哪些物体。其了在工业4.0的时代实现灵活自动化。')
# data2 = get_data('今年3月，受新冠肺炎疫情影响，两会虽没有如期召开')
data.append(data1)
# data.append(data2)

# 获取每句话的单词数量
base_shape = [[len(c) for c in data]]

# 生成预测数据
tensor_words = fluid.create_lod_tensor(data, base_shape, place)

# 执行预测
result = exe.run(program=infer_program,
                 feed={feeded_var_names[0]: tensor_words},
                 fetch_list=target_var)

# 分类名称
names = ['产品', '案例', '解决方案', "爬取错误", "其他内容"]

# 获取结果概率最大的label
for i in range(len(data)):
    lab = np.argsort(result)[0][i][-1]
    print('预测结果标签为：%d， 名称为：%s， 概率为：%f' % (lab, names[lab], result[0][i][lab]))
